
import os,time
# 设置主显示屏
os.environ['DISPLAY'] = ':0'

import cv2 as cv
import numpy as np
import threading
from Drivers import camera
from ComputerVision import commonProcess,searchLaserPoint,searchPath
from Drivers.FOCMotorController import FOCMotorController
from Algorithm.KalmanFilter2D import KalmanFilter2D
from Algorithm.PID import PID


yawMotor = FOCMotorController('/dev/serial/by-id/usb-STMicroelectronics_QDrive_ComPort_207B37B05641-if00')
pinchMotor = FOCMotorController('/dev/serial/by-id/usb-STMicroelectronics_QDrive_ComPort_207A377B5641-if00')

globalLock = threading.Lock()

YAW_OFFSET = 5.445
PINCH_OFFSET = 4.39

YAW_OFFSET = 5.39
PINCH_OFFSET = 4.36

YAW_EDGE_POINT_LIST =    [+5.61,+5.16,+5.16,+5.61,+5.61]
PINCH_EDGE_POINT_LIST =  [+4.14,+4.14,+4.58,+4.58,+4.14]

TASK_RUNNING = False

def drawPoint(
    img: np.ndarray,
    point: list,
    color: list = [255,0,0],
    radius: int = 5
    ):
    img = cv.circle(img,np.int32(point),radius,color,-1)
    return img

def testSearchLaser():
    yawMotor.disable()
    pinchMotor.disable()
    kf = KalmanFilter2D(1,0.05)
    camera.openCamera()
    pred_point = [0,0]
    while True:
        image = camera.captureImage()
        image = commonProcess.cropToRequirement(image)
        redLaserPoint = searchLaserPoint.searchRedLaserPoint(image,searchLaserPoint.METHOD_HSV_RANGE2)
        # if not redLaserPoint:
        #     redLaserPoint = searchLaserPoint.searchRedLaserPoint(image,searchLaserPoint.METHOD_DIFF_CHANNEL)
        print(redLaserPoint)
        if redLaserPoint:
            pred_point = kf.predict(np.float32([redLaserPoint]))[0]
            image = drawPoint(image,pred_point)
        else:
            pred_point = kf.predict(np.float32([pred_point]))[0]
            image = drawPoint(image,pred_point,color=[0,255,0])
        greenLaserPoint = searchLaserPoint.searchGreenLaserPoint(image)
        if greenLaserPoint:
            image = drawPoint(image,greenLaserPoint,[0,0,255])
        cv.imshow('searchLaserPoint',image)
        k = cv.waitKey(1)
        if k == ord('q'):
            break
    cv.destroyAllWindows()
def testSearchPath():
    camera.openCamera()
    image = camera.captureImage()
    image = commonProcess.cropToRequirement(image)
    imageToShow = image.copy()
    image = searchPath.otsuThreshould(image)
    innerContour,outerContour = searchPath.searchInnerOuterContour(image)
    trackPoints = searchPath.generateTrackingPath(innerContour,outerContour)
    for i in range(len(trackPoints)):
        imageToShow = drawPoint(imageToShow,trackPoints[i][0])
    while True:
        cv.imshow('testSearchPath',imageToShow)
        cv.imshow('otsuThresh',image)
        k = cv.waitKey(1)
        if k == ord('q'):
            break
    cv.destroyAllWindows()

def setTaskRunningState(state):
    global TASK_RUNNING
    globalLock.acquire()
    TASK_RUNNING = state
    globalLock.release()

def taskReset():
    # 这是第一道题，回到原点
    yawMotor.enable()
    pinchMotor.enable()
    yawMotor.setAngle(YAW_OFFSET)
    pinchMotor.setAngle(PINCH_OFFSET)
    time.sleep(1)

    setTaskRunningState(False)

def task1():
    # 开环跑一遍扫一下铅笔画的边缘
    global TASK_RUNNING
    yawMotor.enable()
    pinchMotor.enable()
    yawMotor.setAngle(YAW_EDGE_POINT_LIST[0])
    pinchMotor.setAngle(PINCH_EDGE_POINT_LIST[0])
    time.sleep(1.5)
    for i in range(4):
        # 最好是进行一个线性的插值
        yawMotor.setAngle(YAW_EDGE_POINT_LIST[i])
        pinchMotor.setAngle(PINCH_EDGE_POINT_LIST[i])
        currentYaw = YAW_EDGE_POINT_LIST[i]
        currentPinch = PINCH_EDGE_POINT_LIST[i]
        yawInc = (YAW_EDGE_POINT_LIST[i+1] - YAW_EDGE_POINT_LIST[i])/40
        pinchInc = (PINCH_EDGE_POINT_LIST[i+1] - PINCH_EDGE_POINT_LIST[i])/40
        for i in range(40):
            currentYaw += yawInc
            currentPinch += pinchInc
            if TASK_RUNNING == True:
                yawMotor.setAngle(currentYaw)
                pinchMotor.setAngle(currentPinch)
            else:
                return
            time.sleep(0.02)
    yawMotor.setSpeed(0,lowSpeedMode=True)
    pinchMotor.setSpeed(0,lowSpeedMode=True)
    setTaskRunningState(False)

def task2():
    global TASK_RUNNING
    # yawMotor.setAngle(YAW_OFFSET)
    # pinchMotor.setAngle(PINCH_OFFSET)
    yawMotor.setSpeed(0,lowSpeedMode=True)
    pinchMotor.setSpeed(0,lowSpeedMode=True)
    yawMotor.enable()
    pinchMotor.enable()
    # 卡尔曼滤波器
    kf = KalmanFilter2D(1,0.05)
    # 扫描整个矩形框的边缘
    camera.openCamera()
    # 先得到我们要追踪的目标点
    image = camera.captureImage()
    image = commonProcess.cropToRequirement(image)
    imageToShow = image.copy()
    image = searchPath.otsuThreshould(image)
    innerContour,outerContour = searchPath.searchInnerOuterContour(image)
    trackPoints = searchPath.generateTrackingPath(innerContour,outerContour)
    # 得到目标点之后对点之间进行插值操作（否则会出现速度突变）
    interpTargetPoints = None
    for i in range(len(trackPoints)):
        startPoint = trackPoints[i][0]
        stopPoint = trackPoints[(i+1)%len(trackPoints)][0]
        # print(startPoint,stopPoint)
        # 大概插7个点就行
        interpNum = 15
        xInterp = np.linspace(startPoint[0],stopPoint[0],interpNum+1)[:-1]
        yInterp = np.linspace(startPoint[1],stopPoint[1],interpNum+1)[:-1]
        if type(interpTargetPoints) == type(None):
            interpTargetPoints = np.array([xInterp,yInterp]).T
        else:
            interpTargetPoints = np.concatenate((interpTargetPoints,np.array([xInterp,yInterp]).T),axis=0)
    interpTargetPoints = np.concatenate((interpTargetPoints,np.array([interpTargetPoints[0]])))
    yawPID = PID(+0.02,0.000,0,1)
    pinchPID = PID(-0.02,-0.000,0,1)
    MAX_SPEED = 2
    predPoint = [0,0]
    for i in range(len(interpTargetPoints)):
        # 逐个追踪目标点
        targetPoint = interpTargetPoints[i]
        print('追踪点:(%s,%s)'%(targetPoint[0],targetPoint[1]))
        while True:
            if TASK_RUNNING == False:
                globalLock.acquire()
                yawMotor.setSpeed(0,lowSpeedMode=True)
                pinchMotor.setSpeed(0,lowSpeedMode=True)
                globalLock.release()
                return
            else:
                image = camera.captureImage()
                image = commonProcess.cropToRequirement(image)
                # 两种方法混合寻找激光点
                laserPoint = searchLaserPoint.searchRedLaserPoint(image,searchLaserPoint.METHOD_HSV_RANGE2)
                if not laserPoint:
                    laserPoint = searchLaserPoint.searchRedLaserPoint(image,searchLaserPoint.METHOD_HSV_RANGE)
                if not laserPoint:
                    print('未检测到激光点')
                    # 未检测到激光点使用卡尔曼滤波的结果
                    predPoint = kf.predict(np.array([predPoint]))[0]
                else:
                    # 检测到激光点再做追踪
                    predPoint = kf.predict(np.array([laserPoint]))[0]
                    xLoss = predPoint[0] - targetPoint[0]
                    yLoss = predPoint[1] - targetPoint[1]
                    print('激光点坐标:(%s,%s)'%(laserPoint[0],laserPoint[1]))
                    print('相对误差:(%s,%s)'%(xLoss,yLoss))
                    yawFeedback = yawPID.forward(xLoss,0)
                    pinchFeedback = pinchPID.forward(yLoss,0)
                    if yawFeedback > MAX_SPEED: yawFeedback = MAX_SPEED
                    elif yawFeedback < -MAX_SPEED: yawFeedback = -MAX_SPEED
                    if pinchFeedback > MAX_SPEED: pinchFeedback = MAX_SPEED
                    elif pinchFeedback < -MAX_SPEED: pinchFeedback = -MAX_SPEED
                    globalLock.acquire()
                    yawMotor.setSpeed(yawFeedback,lowSpeedMode=True)
                    pinchMotor.setSpeed(pinchFeedback,lowSpeedMode=True)
                    globalLock.release()
                    # 误差小于一定大小就可以找下一个点了
                    if searchPath.relativeDistance([xLoss,yLoss],[0,0]) < 10:
                        break
            
    yawMotor.setSpeed(0,lowSpeedMode=True)
    pinchMotor.setSpeed(0,lowSpeedMode=True)
    setTaskRunningState(False)

    # for i in range(len(interpTargetPoints)):
    #     imageToShow = drawPoint(imageToShow,interpTargetPoints[i])
    # while True:
    #     cv.imshow('interpPoints',imageToShow)
    #     k = cv.waitKey(1)
    #     if k == ord('q'):
    #         break

def runTaskReset():
    global TASK_RUNNING
    if TASK_RUNNING == True:
        print('有任务正在运行')
        return
    else:
        setTaskRunningState(True)
        t = threading.Thread(target=taskReset)
        t.start()

def runTask1():
    global TASK_RUNNING
    if TASK_RUNNING == True:
        print('有任务正在运行')
        return
    else:
        setTaskRunningState(True)
        t = threading.Thread(target=task1)
        t.start()

def runTask2():
    global TASK_RUNNING
    if TASK_RUNNING == True:
        print('有任务正在运行')
        return
    else:
        setTaskRunningState(True)
        t = threading.Thread(target=task2)
        t.start()

def taskEmergeStop():
    # 急停
    global TASK_RUNNING
    globalLock.acquire()
    TASK_RUNNING = False
    yawMotor.setSpeed(0,lowSpeedMode=True)
    pinchMotor.setSpeed(0,lowSpeedMode=True)
    globalLock.release()
    
if __name__ == '__main__':
    
    print('===============')
    print('2023电赛E题1号机程序')
    while True:
        print('(1)复位回原点\n(2)任务1:扫描边框一圈\n(3)任务2:扫描黑色胶带一圈\n(6)使能电机\n(7)禁用电机\n(8)测试:边框寻找功能\n(9)测试:激光定位功能\n(0)退出')
        code = input('请输入任务编号>')
        if code == '1':
            taskReset()
        elif code == '2':
            task1()
        elif code == '3':
            task2()
        elif code == '6':
            yawMotor.enable()
            pinchMotor.enable()
        elif code == '7':
            yawMotor.disable()
            pinchMotor.disable()
        elif code == '8':
            testSearchPath()
        elif code == '9':
            testSearchLaser()
        elif code == '0':
            break
    yawMotor.disable()
    pinchMotor.disable()